#include <libff/common/default_types/ec_pp.hpp>
#include <libsnark/zk_proof_systems/ppzksnark/r1cs_ppzksnark/examples/run_r1cs_ppzksnark.hpp>
#include <iostream>
#include <libsnark/gadgetlib1/protoboard.hpp>
#include <chrono>
#include <sstream>
#include <iomanip>
#include <random>
#include <fstream>
#include <vector>
#include <tuple>
#include <cstdlib>
#include <random>
#include <algorithm>

using namespace libsnark;
using namespace libff;

using clk = std::chrono::high_resolution_clock;
using ms = std::chrono::duration<double, std::milli>;

// Default experiment parameters (can be adjusted)
const size_t DEFAULT_NUM_VEHICLES = 50;
const size_t DEFAULT_ITERATIONS = 100;
const size_t DEFAULT_REPEATS = 10;

// Network simulation defaults (ms)
const double DEFAULT_NETWORK_DELAY_MS = 20.0; // per-message one-way delay
const double NETWORK_JITTER_MS = 5.0;

// Helper: build a very small R1CS (x * y = z)
r1cs_constraint_system<Fr<default_ec_pp>> make_simple_cs()
{
    protoboard<Fr<default_ec_pp>> pb;
    pb_variable<Fr<default_ec_pp>> x, y, z;
    x.allocate(pb, "x");
    y.allocate(pb, "y");
    z.allocate(pb, "z");
    pb.set_input_sizes(1);
    pb.add_r1cs_constraint(r1cs_constraint<Fr<default_ec_pp>>(x, y, z));
    pb.val(x) = 3;
    pb.val(y) = 4;
    pb.val(z) = 12;
    return pb.get_constraint_system();
}

struct Stats {
    double avg_prove_ms = 0;
    double avg_verify_ms = 0;
    double avg_proof_bytes = 0;
    size_t total_comm_bytes = 0;
};

// Simulate vehicle-to-vehicle authentication
Stats simulate_v2v(size_t num_vehicles, size_t iterations, bool vk_pre_distributed, bool use_processed_vk,
                   std::ofstream* csv /*= nullptr*/, std::ofstream* rounds_csv = nullptr,
                   const std::string& scenario_name = "", size_t repeat_idx = 0)
{
    std::cout << "\n== V2V simulation: vehicles=" << num_vehicles << ", iter=" << iterations
              << ", vk_pre_distributed=" << vk_pre_distributed << ", processed_vk=" << use_processed_vk << " ==\n";

    auto cs = make_simple_cs();

    // Pre-generate keypairs for vehicles
    std::vector<r1cs_ppzksnark_keypair<default_ec_pp>> keys;
    keys.reserve(num_vehicles);
    for (size_t i = 0; i < num_vehicles; ++i) {
        keys.emplace_back(r1cs_ppzksnark_generator<default_ec_pp>(cs));
    }

    // Pre-serialize VK/PVK if pre-distributed
    std::vector<std::string> vk_ser(num_vehicles);
    std::vector<std::string> pvk_ser(num_vehicles);
    std::vector<r1cs_ppzksnark_processed_verification_key<default_ec_pp>> pvks(num_vehicles);
    for (size_t i = 0; i < num_vehicles; ++i) {
        std::ostringstream vkoss; vkoss << keys[i].vk; vk_ser[i] = vkoss.str();
        if (use_processed_vk) {
            pvks[i] = r1cs_ppzksnark_verifier_process_vk<default_ec_pp>(keys[i].vk);
            std::ostringstream pvkoss; pvkoss << pvks[i]; pvk_ser[i] = pvkoss.str();
        }
    }

    std::mt19937_64 rng(12345);
    std::uniform_int_distribution<size_t> uniform(0, num_vehicles - 1);
    std::normal_distribution<double> net_delay(DEFAULT_NETWORK_DELAY_MS, NETWORK_JITTER_MS);

    double total_prove_ms = 0;
    double total_verify_ms = 0;
    size_t total_proof_bytes = 0;
    size_t total_comm_bytes = 0;

    // inputs (same for all runs)
    protoboard<Fr<default_ec_pp>> pb;
    pb_variable<Fr<default_ec_pp>> x, y, z;
    x.allocate(pb, "x"); y.allocate(pb, "y"); z.allocate(pb, "z"); pb.set_input_sizes(1);
    pb.val(x) = 3; pb.val(y) = 4; pb.val(z) = 12;
    const auto primary_input = pb.primary_input();
    const auto auxiliary_input = pb.auxiliary_input();

    // Precompute one proof per key to avoid repeated expensive prover calls.
    std::vector<r1cs_ppzksnark_proof<default_ec_pp>> proofs(num_vehicles);
    std::vector<std::string> proof_ser(num_vehicles);
    std::vector<size_t> proof_size(num_vehicles);
    for (size_t i = 0; i < num_vehicles; ++i) {
        proofs[i] = r1cs_ppzksnark_prover<default_ec_pp>(keys[i].pk, primary_input, auxiliary_input);
        std::ostringstream poss; poss << proofs[i]; proof_ser[i] = poss.str(); proof_size[i] = proof_ser[i].size();
    }

    // Pre-deserialize VK/PVK if they will be sent on-demand so we don't pay per-iteration
    // deserialization cost. Keep serialized sizes for comm accounting.
    std::vector<r1cs_ppzksnark_verification_key<default_ec_pp>> vk_deser;
    std::vector<r1cs_ppzksnark_processed_verification_key<default_ec_pp>> pvk_deser;
    if (!vk_pre_distributed) {
        if (use_processed_vk) {
            pvk_deser.resize(num_vehicles);
            for (size_t i = 0; i < num_vehicles; ++i) { std::istringstream iss(pvk_ser[i]); iss >> pvk_deser[i]; }
        } else {
            vk_deser.resize(num_vehicles);
            for (size_t i = 0; i < num_vehicles; ++i) { std::istringstream iss(vk_ser[i]); iss >> vk_deser[i]; }
        }
    }

    for (size_t it = 0; it < iterations; ++it) {
        size_t p = uniform(rng);
        size_t v = uniform(rng);
        if (v == p) v = (v + 1) % num_vehicles;

    // Prover p generates proof (reuse cached proof to reduce cost)
    auto t1 = clk::now();
    // Simulate round 0: Verifier -> Prover (nonce/challenge), small fixed size, network delay
        double r0_delay = std::max(1.0, net_delay(rng));
        size_t r0_bytes = 32; // nonce/challenge size
        if (rounds_csv) {
            size_t vk_b_r0 = 0; // verifier->prover round does not send VK
            size_t verifier_recv_r0 = (vk_pre_distributed ? (use_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : 0);
            (*rounds_csv) << scenario_name << ",V2V," << repeat_idx << "," << it << ",0,verifier,prover," << r0_bytes << "," << std::fixed << std::setprecision(3) << r0_delay << "," << vk_b_r0 << "," << verifier_recv_r0 << "\n";
        }
        // account for network delay before prover can start
        // (we won't actually sleep; just include delay in round-level accounting)

    // use cached proof
    auto &proof = proofs[p];
    auto t2 = clk::now();
    double prove_ms = ms(t2 - t1).count();
    size_t proof_b = proof_size[p];

        // communication: proof + optional vk/pvk
        size_t comm = proof_b;
        if (!vk_pre_distributed) {
            if (use_processed_vk) comm += pvk_ser[p].size(); else comm += vk_ser[p].size();
        }

        // Simulate round 1: Prover -> Verifier (proof [+ vk/pvk if on-demand])
        double r1_delay = std::max(1.0, net_delay(rng));
        size_t r1_bytes = proof_b;
        if (!vk_pre_distributed) { r1_bytes += (use_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()); }
        if (rounds_csv) {
            size_t vk_b_r1 = (!vk_pre_distributed ? (use_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : 0);
            size_t verifier_recv_r1 = (vk_pre_distributed ? (use_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : vk_b_r1);
            (*rounds_csv) << scenario_name << ",V2V," << repeat_idx << "," << it << ",1,prover,verifier," << r1_bytes << "," << std::fixed << std::setprecision(3) << r1_delay << "," << vk_b_r1 << "," << verifier_recv_r1 << "\n";
        }

        // Verifier verifies (use cached deserialized vk/pvk when we preloaded them)
        auto t3 = clk::now();
        bool ok;
        if (use_processed_vk) {
            if (vk_pre_distributed) {
                ok = r1cs_ppzksnark_online_verifier_strong_IC<default_ec_pp>(pvks[p], primary_input, proof);
            } else {
                // use cached deserialized pvk
                ok = r1cs_ppzksnark_online_verifier_strong_IC<default_ec_pp>(pvk_deser[p], primary_input, proof);
            }
        } else {
            if (vk_pre_distributed) {
                ok = r1cs_ppzksnark_verifier_strong_IC<default_ec_pp>(keys[p].vk, primary_input, proof);
            } else {
                // use cached deserialized vk
                ok = r1cs_ppzksnark_verifier_strong_IC<default_ec_pp>(vk_deser[p], primary_input, proof);
            }
        }
        auto t4 = clk::now();
        double verify_ms = ms(t4 - t3).count();

        total_prove_ms += prove_ms;
        total_verify_ms += verify_ms;
        total_proof_bytes += proof_b;
        total_comm_bytes += comm;

        // CSV row now includes vk_bytes and verifier_recv_bytes
        if (csv) {
            size_t vk_b = 0;
            if (!vk_pre_distributed) vk_b = (use_processed_vk ? pvk_ser[p].size() : vk_ser[p].size());
            // verifier_recv_bytes: how many bytes the verifier has already (pre-distributed) or receives this round
            size_t verifier_recv = (vk_pre_distributed ? (use_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : vk_b);
            (*csv) << scenario_name << ",V2V," << repeat_idx << "," << it << "," << p << "," << v << ","
                   << std::fixed << std::setprecision(6) << prove_ms << "," << verify_ms << "," << proof_b << ","
                   << vk_b << "," << verifier_recv << "," << comm << "\n";
        }
    }

    Stats s;
    s.avg_prove_ms = total_prove_ms / iterations;
    s.avg_verify_ms = total_verify_ms / iterations;
    s.avg_proof_bytes = double(total_proof_bytes) / iterations;
    s.total_comm_bytes = total_comm_bytes;

    std::cout << "V2V avg prove(ms): " << s.avg_prove_ms << ", avg verify(ms): " << s.avg_verify_ms << "\n";
    std::cout << "V2V avg proof bytes: " << s.avg_proof_bytes << ", total comm bytes: " << s.total_comm_bytes << "\n";
    return s;
}

// Simulate vehicle-to-RSU authentication (RSU is single verifier)
Stats simulate_v2i(size_t num_vehicles, size_t iterations, bool vk_pre_distributed_to_rsu, bool rsu_uses_processed_vk,
                   std::ofstream* csv /*= nullptr*/, std::ofstream* rounds_csv = nullptr,
                   const std::string& scenario_name = "", size_t repeat_idx = 0)
{
    std::cout << "\n== V2I simulation: vehicles=" << num_vehicles << ", iter=" << iterations
              << ", rsu_has_vk=" << vk_pre_distributed_to_rsu << ", rsu_processed_vk=" << rsu_uses_processed_vk << " ==\n";

    auto cs = make_simple_cs();
    std::vector<r1cs_ppzksnark_keypair<default_ec_pp>> keys;
    keys.reserve(num_vehicles);
    for (size_t i = 0; i < num_vehicles; ++i) keys.emplace_back(r1cs_ppzksnark_generator<default_ec_pp>(cs));

    // RSU pre-stores VK/PVK for vehicles or will receive on-demand
    std::vector<std::string> vk_ser(num_vehicles);
    std::vector<std::string> pvk_ser(num_vehicles);
    std::vector<r1cs_ppzksnark_processed_verification_key<default_ec_pp>> pvks(num_vehicles);
    for (size_t i = 0; i < num_vehicles; ++i) {
        std::ostringstream vkoss; vkoss << keys[i].vk; vk_ser[i] = vkoss.str();
        if (rsu_uses_processed_vk) { pvks[i] = r1cs_ppzksnark_verifier_process_vk<default_ec_pp>(keys[i].vk); std::ostringstream pss; pss << pvks[i]; pvk_ser[i] = pss.str(); }
    }

    std::mt19937_64 rng(123);
    std::uniform_int_distribution<size_t> uniform(0, num_vehicles - 1);
    std::normal_distribution<double> net_delay(DEFAULT_NETWORK_DELAY_MS, NETWORK_JITTER_MS);

    double total_prove_ms = 0; double total_verify_ms = 0; size_t total_proof_bytes = 0; size_t total_comm_bytes = 0;

    // inputs
    protoboard<Fr<default_ec_pp>> pb;
    pb_variable<Fr<default_ec_pp>> x, y, z;
    x.allocate(pb, "x"); y.allocate(pb, "y"); z.allocate(pb, "z"); pb.set_input_sizes(1);
    pb.val(x) = 3; pb.val(y) = 4; pb.val(z) = 12;
    const auto primary_input = pb.primary_input();
    const auto auxiliary_input = pb.auxiliary_input();
    // Precompute one proof per key to avoid repeated expensive prover calls.
    std::vector<r1cs_ppzksnark_proof<default_ec_pp>> proofs(num_vehicles);
    std::vector<std::string> proof_ser(num_vehicles);
    std::vector<size_t> proof_size(num_vehicles);
    for (size_t i = 0; i < num_vehicles; ++i) {
        proofs[i] = r1cs_ppzksnark_prover<default_ec_pp>(keys[i].pk, primary_input, auxiliary_input);
        std::ostringstream poss; poss << proofs[i]; proof_ser[i] = poss.str(); proof_size[i] = proof_ser[i].size();
    }

    // Pre-deserialize VK/PVK if they will be sent on-demand so we don't pay per-iteration
    // deserialization cost. Keep serialized sizes for comm accounting.
    std::vector<r1cs_ppzksnark_verification_key<default_ec_pp>> vk_deser;
    std::vector<r1cs_ppzksnark_processed_verification_key<default_ec_pp>> pvk_deser;
    if (!vk_pre_distributed_to_rsu) {
        if (rsu_uses_processed_vk) {
            pvk_deser.resize(num_vehicles);
            for (size_t i = 0; i < num_vehicles; ++i) { std::istringstream iss(pvk_ser[i]); iss >> pvk_deser[i]; }
        } else {
            vk_deser.resize(num_vehicles);
            for (size_t i = 0; i < num_vehicles; ++i) { std::istringstream iss(vk_ser[i]); iss >> vk_deser[i]; }
        }
    }

    for (size_t it = 0; it < iterations; ++it) {
        size_t p = uniform(rng);

        // Round 0: RSU -> Vehicle (nonce/challenge)
        double r0_delay = std::max(1.0, net_delay(rng));
        size_t r0_bytes = 32;
        if (rounds_csv) {
            size_t vk_b_r0 = 0;
            size_t verifier_recv_r0 = (vk_pre_distributed_to_rsu ? (rsu_uses_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : 0);
            (*rounds_csv) << scenario_name << ",V2I," << repeat_idx << "," << it << ",0,rsu,vehicle," << r0_bytes << "," << std::fixed << std::setprecision(3) << r0_delay << "," << vk_b_r0 << "," << verifier_recv_r0 << "\n";
        }

    auto t1 = clk::now();
    // reuse cached proof
    auto &proof = proofs[p];
    auto t2 = clk::now();
    double prove_ms = ms(t2 - t1).count();
    size_t proof_b = proof_size[p];

        size_t comm = proof_b;
        if (!vk_pre_distributed_to_rsu) {
            if (rsu_uses_processed_vk) comm += pvk_ser[p].size(); else comm += vk_ser[p].size();
        }

        // Round 1: Vehicle -> RSU (proof [+ vk/pvk if on-demand])
        double r1_delay = std::max(1.0, net_delay(rng));
        size_t r1_bytes = proof_b;
        if (!vk_pre_distributed_to_rsu) { r1_bytes += (rsu_uses_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()); }
        if (rounds_csv) {
            size_t vk_b_r1 = (!vk_pre_distributed_to_rsu ? (rsu_uses_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : 0);
            size_t verifier_recv_r1 = (vk_pre_distributed_to_rsu ? (rsu_uses_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : vk_b_r1);
            (*rounds_csv) << scenario_name << ",V2I," << repeat_idx << "," << it << ",1,vehicle,rsu," << r1_bytes << "," << std::fixed << std::setprecision(3) << r1_delay << "," << vk_b_r1 << "," << verifier_recv_r1 << "\n";
        }

        auto t3 = clk::now();
        bool ok;
        if (rsu_uses_processed_vk) {
            if (vk_pre_distributed_to_rsu) ok = r1cs_ppzksnark_online_verifier_strong_IC<default_ec_pp>(pvks[p], primary_input, proof);
            else ok = r1cs_ppzksnark_online_verifier_strong_IC<default_ec_pp>(pvk_deser[p], primary_input, proof);
        } else {
            if (vk_pre_distributed_to_rsu) ok = r1cs_ppzksnark_verifier_strong_IC<default_ec_pp>(keys[p].vk, primary_input, proof);
            else ok = r1cs_ppzksnark_verifier_strong_IC<default_ec_pp>(vk_deser[p], primary_input, proof);
        }
        auto t4 = clk::now(); double verify_ms = ms(t4 - t3).count();

        total_prove_ms += prove_ms; total_verify_ms += verify_ms; total_proof_bytes += proof_b; total_comm_bytes += comm;

        // CSV row now includes vk_bytes and verifier_recv_bytes
        if (csv) {
            size_t vk_b = 0;
            if (!vk_pre_distributed_to_rsu) vk_b = (rsu_uses_processed_vk ? pvk_ser[p].size() : vk_ser[p].size());
            size_t verifier_recv = (vk_pre_distributed_to_rsu ? (rsu_uses_processed_vk ? pvk_ser[p].size() : vk_ser[p].size()) : vk_b);
            (*csv) << scenario_name << ",V2I," << repeat_idx << "," << it << "," << p << ",," 
                   << std::fixed << std::setprecision(6) << prove_ms << "," << verify_ms << "," << proof_b << ","
                   << vk_b << "," << verifier_recv << "," << comm << "\n";
        }
    }

    Stats s; s.avg_prove_ms = total_prove_ms / iterations; s.avg_verify_ms = total_verify_ms / iterations; s.avg_proof_bytes = double(total_proof_bytes) / iterations; s.total_comm_bytes = total_comm_bytes;
    std::cout << "V2I avg prove(ms): " << s.avg_prove_ms << ", avg verify(ms): " << s.avg_verify_ms << "\n";
    std::cout << "V2I avg proof bytes: " << s.avg_proof_bytes << ", total comm bytes: " << s.total_comm_bytes << "\n";
    return s;
}

int main()
{
    default_ec_pp::init_public_params();
    std::cout << "✅ BN128 curve initialized" << std::endl;

    size_t num_vehicles = DEFAULT_NUM_VEHICLES;
    size_t iterations = DEFAULT_ITERATIONS;
    size_t repeats = DEFAULT_REPEATS;

    const std::string base_dir = std::string("/home/") + std::string(std::getenv("USER") ? std::getenv("USER") : "user") + "/zk_demo/build";
    const std::string csv_path = base_dir + "/results.csv";
    const std::string rounds_csv_path = base_dir + "/results_rounds.csv";
    std::ofstream csv(csv_path, std::ios::out | std::ios::trunc);
    std::ofstream rounds_csv(rounds_csv_path, std::ios::out | std::ios::trunc);
    if (!csv) {
        std::cerr << "Failed to open CSV output: " << csv_path << std::endl;
        return 1;
    }
    if (!rounds_csv) {
        std::cerr << "Failed to open rounds CSV output: " << rounds_csv_path << std::endl;
        return 1;
    }
    csv << "scenario,type,repeat,iter,prover,verifier,prove_ms,verify_ms,proof_bytes,vk_bytes,verifier_recv_bytes,comm_bytes\n";
    rounds_csv << "scenario,type,repeat,iter,round,from,to,sent_bytes,latency_ms,vk_bytes,verifier_recv_bytes\n";

    // allow overrides via environment variables for quick testing
    if (std::getenv("NUM_VEHICLES")) num_vehicles = std::stoul(std::getenv("NUM_VEHICLES"));
    if (std::getenv("ITERATIONS")) iterations = std::stoul(std::getenv("ITERATIONS"));
    if (std::getenv("REPEATS")) repeats = std::stoul(std::getenv("REPEATS"));

    for (size_t r = 0; r < repeats; ++r) {
        simulate_v2v(num_vehicles, iterations, true, false, &csv, &rounds_csv, "v2v_vk_pre_nonproc", r);
        simulate_v2v(num_vehicles, iterations, false, false, &csv, &rounds_csv, "v2v_vk_on_demand", r);
        simulate_v2v(num_vehicles, iterations, true, true, &csv, &rounds_csv, "v2v_vk_pre_proc", r);
        simulate_v2i(num_vehicles, iterations, true, true, &csv, &rounds_csv, "v2i_rsu_has_pvk", r);
        simulate_v2i(num_vehicles, iterations, false, true, &csv, &rounds_csv, "v2i_rsu_gets_pvk", r);
    }

    csv.close();
    std::cout << "CSV results written to: " << csv_path << std::endl;

    return 0;
}

